# -*- coding: utf-8 -*-
# time: 2025/4/17 15:30
# file: mp3_to_text.py.py
# author: hanson
import torch
import whisper
from transformers import pipeline
from typing import Tuple
import warnings

warnings.filterwarnings("ignore")  # 忽略librosa的警告


def mp3_to_text(
        audio_path: str,
        asr_model: str = "tiny",
        summarize: bool = False
) -> Tuple[str, str]:
    """
    MP3转文本 + 内容摘要（可选）
    :param audio_path: MP3文件路径
    :param asr_model: Whisper模型大小 (tiny/base/small)
    :param summarize: 是否生成摘要
    :return: (转录文本, 摘要文本)
    """
    # 1. 语音识别（Whisper）
   # model = whisper.load_model(asr_model, device="cpu")

    # 使用VITS轻量版（需额外安装）
    from transformers import AutoModelForSpeechSeq2Seq
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        "SUPL23/VITS-zh-25M",
        torch_dtype=torch.float32,
        device_map="cpu"
    )
    result = model.transcribe(audio_path, language="zh" if "zh" in audio_path else "en")
    transcription = result["text"]

    # 2. 内容理解（DistilBERT）
    summary = ""
    if summarize and len(transcription) > 10:
        summarizer = pipeline(
            "summarization",
            model="distilbart-cnn-6-6",
            framework="pt",
            device=-1  # 强制使用CPU
        )
        summary = summarizer(transcription, max_length=50, min_length=10)[0]["summary_text"]

    return transcription, summary


if __name__ == "__main__":
    import sys

    audio_file = sys.argv[1] if len(sys.argv) > 1 else "./mp3/close.mp3"

    print("===== 原始转录 =====")
    text, summary = mp3_to_text(audio_file, summarize=True)
    print(text)

    print("\n===== 内容摘要 =====")
    print(summary if summary else "（文本过短无需摘要）")